/* pirateAcorn.c
 *
 * Used to map the control flow of read and write operations in the Linux kernel
 *
 * Definitions:
 * Acorn   -Slang for someone who is performing an intelligence function
 * Pirate  -Totally cooler than Ninjas
 *
 * (C) Copyright 2008 by Matt Weir, (weir@cs.fsu.edu)
 *
 * This file is released under the GPLv2.
 */

#include <linux/sched.h>
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/marker.h>
#include <asm/atomic.h>
#include <linux/moduleparam.h>
#include <linux/init.h>
#include <linux/fs.h>
#include <asm/uaccess.h>

#include "pirateAcorn.h"

struct probe_data {
	const char *name;
	const char *format;
	marker_probe_func *probe_func;
};


struct probe_manager_struct {
        struct semaphore sem;           //  mutual exclusion semaphore
	int isProbeRead;		//  if it is currently probing read requests, used to unregister the marker afterwards
	int isProbeWrite;		//  if it is currently probing write requests, used to unregister the marker afterwards
	int isSingleMonitor;		//  if it is only monitering a single thread
	int isFinished;			//  used during the exit procedure to keep new logging from starting after it gives up the sem
};

static int logger_pid=0;                //  The pid of syslog or its equvalent, most of the time you won't want to view these logs
module_param(logger_pid, int, S_IRUGO); //  can be specified at mod init.

static int major = 0;   /* dynamic by default */
module_param(major, int, 0);

static struct probe_manager_struct probeManager;	//  Holds most of the record keeping and state info

/////////////////////////////////////////////////////////////////
//Function: probe_subsystem_vfs_read
//
//Used to collect data from the vfs read events
//Yes, right now it is similar(the same) as probe_subsystem_vfs_write
//but I wanted to seperate them in case I ever want to change how
//they behave.
//
//Right now it collects 
//--the function the maker was called
//--the pid of the calling thread
//--the time the marker was called
//--the size of the read
//
//Output
//--currently just does a prink to the log
//--could be extended later to use some other data logging method
//
void probe_subsystem_vfs_read(void *probe_data, void *call_data,
	const char *format, va_list *args)
{	
	/* Declare args */
       
	const char *print_name; 
        int print_pid;
	long long print_time;
	unsigned int print_size;
	if (down_interruptible(&probeManager.sem)) {
 	       return;
        }
//-->locked
	if (!probeManager.isFinished) { //there were some cases where it was still trying to print after removal
		/* Assign args */
		print_name = va_arg(*args,typeof(print_name));
		print_pid = va_arg(*args, typeof(print_pid));
        	print_time = va_arg(*args, typeof(print_time));
		print_size = va_arg(*args, typeof(print_size));
		va_end(*args);   //clean up for va_list

		/* Call printk */
		printk(KERN_INFO "PIRATEACORN: %s: Pid=%i Time=%lli Size of read= %u\n", print_name, print_pid, print_time, print_size);
	}
	up(&probeManager.sem);
//<--unlocked
}


/////////////////////////////////////////////////////////////////
//Function: probe_subsystem_vfs_write
//
//Used to collect data from the vfs write events
//Yes, right now it is similar(the same) as probe_subsystem_vfs_read
//but I wanted to seperate them in case I ever want to change how
//they behave.
//
//Right now it collects 
//--the function the marker was called
//--the pid of the calling thread
//--the time the marker was called
//--the size of the read
//
//Output
//--currently just does a prink to the log
//--could be extended later to use some other data logging method
//
void probe_subsystem_vfs_write(void *probe_data, void *call_data,
        const char *format, va_list *args)
{
        /* Declare args */

        const char *print_name;
        int print_pid;
        long long print_time;
        unsigned int print_size;
	if (down_interruptible(&probeManager.sem)) {
        	return;
        }
//-->locked
	if (!probeManager.isFinished) {
        	/* Assign args */
        	print_name = va_arg(*args,typeof(print_name));
        	print_pid = va_arg(*args, typeof(print_pid));
       		print_time = va_arg(*args, typeof(print_time));
        	print_size = va_arg(*args, typeof(print_size));
        	va_end(*args);   //clean up for va_list

        	/* Call printk */
        	printk(KERN_INFO "PIRATEACORN: %s: Pid=%i Time=%lli Size of write= %u\n", print_name, print_pid, print_time, print_size);
	}
	up(&probeManager.sem);
//<--unlocked
	return;
}



//////////////////////////////////////////////////////////
//This structure holds all the read probes
//When it is registered, all read markers will start fireing
//Note, if you only want certain read makers to fire,
//you can edit/duplicate this structure with only the
//markers you want to use.
//
static struct probe_data probe_readAll_array[] =
{
	{	.name = "sys_read_start",
		.format = "%s %i %lli %u",
		.probe_func = probe_subsystem_vfs_read },
	{       .name = "sys_read_end",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "vfs_read_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "do_sync_read_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "generic_file_aio_read_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "generic_file_direct_IO_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "sys_pread64_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "sys_pread64_stop",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "sys_generic_file_read_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "do_sync_read_end",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "do_sync_read_forloop",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "vfs_read_fop_read",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "vfs_read_fop_aioread",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "seq_read_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },
	{       .name = "relay_file_read_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_read },

};


//////////////////////////////////////////////////////////
//This structure holds all the write probes
//When it is registered, all write markers will start fireing
//Note, if you only want certain write makers to fire,
//you can edit/duplicate this structure with only the
//markers you want to use.
//
static struct probe_data probe_writeAll_array[] =
{
	{       .name = "sys_write_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_write },
	{       .name = "sys_write_stop",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_write },
	{       .name = "vfs_write_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_write },
	{       .name = "do_sync_write_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_write },
	{       .name = "sys_pwrite64_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_write },
	{       .name = "sys_pwrite64_start",
                .format = "%s %i %lli %u",
                .probe_func = probe_subsystem_vfs_write },
};


////////////////////////////////////////////////////////
//The default open function
//Currently doesn't do anything
int pa_open (struct inode *inode, struct file *filp) {
        return 0;
}

////////////////////////////////////////////////////////
//The default release function
//Currently doesn't do anything
int pa_release (struct inode *inode, struct file *filp) {
        return 0;
}

////////////////////////////////////////////////////////
//The default read function
//Currently doesn't do anything
ssize_t pa_read(struct file *filp, char __user *buf, size_t count, loff_t *f_pos) {
        return 0;
}

////////////////////////////////////////////////////////
//The default write function
//Currently doesn't do anything
ssize_t pa_write(struct file *filp, const char __user *buf, size_t count, loff_t *f_pos) {
        return 0;
}




int pa_ioctl(struct inode *inode, struct file *filp,
                 unsigned int cmd, unsigned long arg)
{	
	int i;
	int retVal=0;
	int result;
        /*
         * extract the type and number bitfields, and don't decode
         * wrong cmds: return ENOTTY (inappropriate ioctl) before access_ok()
         */
        if (_IOC_TYPE(cmd) != PA_IOC_MAGIC) return -ENOTTY;
        if (_IOC_NR(cmd) > PA_IOC_MAXNR) return -ENOTTY;
        switch(cmd) {
          case PA_IOCREADSTART: //start markers for reads
		if (down_interruptible(&probeManager.sem)) {
                        return -ERESTARTSYS;                   
        	}
//-->locked
		if ((probeManager.isProbeRead==0)&&(!probeManager.isFinished)) {  //it is not already tracing reads
			retVal=1;
			for (i = 0; i < ARRAY_SIZE(probe_readAll_array); i++) {
                		result = marker_probe_register(probe_readAll_array[i].name, probe_readAll_array[i].format,
					 probe_readAll_array[i].probe_func, &probe_readAll_array[i]);
                		if (result) {
                        		printk(KERN_INFO "PIRATEACORN: Unable to register probe %s\n",probe_readAll_array[i].name);
					retVal=ERESTARTSYS;
				}
			}
			if (retVal==1) {
				probeManager.isProbeRead=1;
			}
        	}
		else { //it is already following reads
			retVal=0;
		}
		up(&probeManager.sem);
//<--unlocked
          	break;
	  case PA_IOCREADSTOP:   //stop tracing reads
                if (down_interruptible(&probeManager.sem)) {
                        return -ERESTARTSYS;
                }
//-->locked
                if (probeManager.isProbeRead==0) {  //it is not already tracing reads
			retVal=1;
		}
		else {
			for (i = 0; i < ARRAY_SIZE(probe_readAll_array); i++) {
                        	marker_probe_unregister(probe_readAll_array[i].name, probe_readAll_array[i].probe_func, &probe_readAll_array[i]);
                	}
			retVal=1;
			probeManager.isProbeRead=0;
		}		
		up(&probeManager.sem);
//<--unlocked
                break;

//------------------------IOC for starting writes-------------------------//
	  case PA_IOCWRITESTART:
                if (down_interruptible(&probeManager.sem)) {
                        return -ERESTARTSYS;
                }
//-->locked
                if ((probeManager.isProbeWrite==0)&&(!probeManager.isFinished)) {  //it is not already tracing writes
                        retVal=1;
                        for (i = 0; i < ARRAY_SIZE(probe_writeAll_array); i++) {
                                result = marker_probe_register(probe_writeAll_array[i].name, probe_writeAll_array[i].format,
                                         probe_writeAll_array[i].probe_func, &probe_writeAll_array[i]);
                                if (result) {
                                        printk(KERN_INFO "PIRATEACORN: Unable to register probe %s\n",probe_writeAll_array[i].name);
                                        retVal=ERESTARTSYS;
                                }
                        }
                        if (retVal==1) {
                                probeManager.isProbeWrite=1;
                        }
                }
                else { //it is already following writes
                        retVal=0;
                }
                up(&probeManager.sem);
//<--unlocked

                break;
          case PA_IOCWRITESTOP:
                if (down_interruptible(&probeManager.sem)) {
                        return -ERESTARTSYS;
                }
//-->locked
                if (probeManager.isProbeWrite==0) {  //it is not already tracing writes
                        retVal=1;
                }
                else {
                        for (i = 0; i < ARRAY_SIZE(probe_writeAll_array); i++) {
                                marker_probe_unregister(probe_writeAll_array[i].name, probe_writeAll_array[i].probe_func, &probe_writeAll_array[i]);
                        }
                        retVal=1;
                        probeManager.isProbeWrite=0;
                }
                up(&probeManager.sem);
//<--unlocked
                break;

	  case PA_IOCLOGGER:
		result=(int)arg;
		if (result>0) {
			marker_set_logger(result);
			printk(KERN_INFO "PIRATEACORN:Set logger to %i\n", result);
		}
		printk(KERN_INFO "PIRATEACORN:Set logger to %i\n", result);
                break;
	  case PA_IOCNOLOGGER:
 		if (down_interruptible(&probeManager.sem)) {
                        return -ERESTARTSYS;
                }
//-->locked
                if (!probeManager.isFinished) {
			marker_clear_logger();
		}
		up(&probeManager.sem);
//<--unlocked

                break;
	  case PA_IOCLOGALL:
                if (down_interruptible(&probeManager.sem)) {
                        return -ERESTARTSYS;
                }
//-->locked
                if (!probeManager.isFinished) {  
			marker_log_all();
		}
                up(&probeManager.sem);
//<--unlocked

                break;
	  case PA_IOCLOGPID:
		result=(int)arg;
                if (result>0) {
                        marker_log_single(result);
                }

                break;
          default:
                return -ENOTTY;
        }
        return 0;
}



struct file_operations pa_fops = {
        .owner   = THIS_MODULE,
        .read    = pa_read,
        .write   = pa_write,
        .open    = pa_open,
        .release = pa_release,
        .ioctl   = pa_ioctl,
};





static int __init probe_init(void)
{
	int result;
	printk(KERN_INFO "PIRATEACORN: STARTING PROBE\n");
	
	result = register_chrdev(major, "pirateAcorn", &pa_fops);
        if (result < 0) {
                printk(KERN_INFO "PIRATEACORN: can't get major number\n");
                return result;
        }
	if (major == 0) major = result; /* dynamic */


        //------Initilize variables, by default is set to monitor everything
        probeManager.isProbeRead=0;
        probeManager.isProbeWrite=0;
	probeManager.isSingleMonitor=0;
	probeManager.isFinished=0;
	init_MUTEX(&probeManager.sem);
        if (logger_pid > 0) { //don't know what the maximum pid can be so I'm not checking that bounding
		if (marker_set_logger(logger_pid)) {
			printk(KERN_ALERT "PIRATEACORN: UNABLE TO TALK TO THE MARKER FUNCTION\n");
			return -EACCES;
		}
	}
/*
	for (i = 0; i < ARRAY_SIZE(probe_array); i++) {
		result = marker_probe_register(probe_array[i].name,
				probe_array[i].format,
				probe_array[i].probe_func, &probe_array[i]);
		if (result)
			printk(KERN_INFO "PIRATEACORN: Unable to register probe %s\n",probe_array[i].name);
	}

*/
	return 0;
}

static void __exit probe_fini(void)
{
	int i;
	//---------Lock so no one turns on a probe while deregistering everything----//
//-->Locking
	if (down_interruptible(&probeManager.sem)) {  //what happens when you fail
                        return;                      //in the cleanup function?
        }
	if (probeManager.isProbeRead) {   //unregister the probe read monitor
		for (i = 0; i < ARRAY_SIZE(probe_readAll_array); i++) {
			marker_probe_unregister(probe_readAll_array[i].name, probe_readAll_array[i].probe_func, &probe_readAll_array[i]);
		}
	}
	if (probeManager.isProbeWrite) {   //unregister the probe write monitor
                for (i = 0; i < ARRAY_SIZE(probe_writeAll_array); i++) {
                        marker_probe_unregister(probe_writeAll_array[i].name, probe_writeAll_array[i].probe_func, &probe_writeAll_array[i]);
                }
        }
        //-----Clears the marker settings back to the default
	marker_log_all();
	marker_clear_logger();

	probeManager.isFinished=1;  //This way no new probes will be registered
	up(&probeManager.sem);      //release lock
//<--Unlocked
	printk(KERN_INFO "PIRATEACORN: Exiting\n");
	unregister_chrdev(major, "pirateAcorn");
	return;
}

module_init(probe_init);
module_exit(probe_fini);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Matt Weir");
MODULE_DESCRIPTION("Read/Write kernel function tracer");
